# -*- coding: utf-8 -*-
"""
Module for post-processing and data output.

This module provides functions to save the simulation results in a format
that can be visualized by external software like ParaView.
"""
import os
import logging
import numpy as np

from .mesh import Mesh
from .solver import CoupledSolver

class PostProcessor:
    """
    Handles the writing of simulation data to files.
    """
    def __init__(self, mesh: Mesh, solver: CoupledSolver, config):
        """
        Initializes the PostProcessor.

        Args:
            mesh (Mesh): The computational mesh object.
            solver (CoupledSolver): The fluid solver object.
            config: The simulation configuration object.
        """
        self.mesh = mesh
        self.solver = solver
        self.config = config
        self.output_dir = self.config.OUTPUT_DIR
        
        # Ensure the output directory exists
        if not os.path.exists(self.output_dir):
            os.makedirs(self.output_dir)
            logging.info(f"Created output directory: {self.output_dir}")

    def write_output(self, time_step_index: int):
        """
        Writes the current flow field data to a VTK file.

        Args:
            time_step_index (int): The index of the current time step.
        """
        file_path = os.path.join(self.output_dir, f"result_{time_step_index:05d}.vtk")
        
        logging.debug(f"Writing output to {file_path}...")
        
        # Prepare data for writing.
        # meshio expects point data as a dictionary of arrays.
        point_data = {
            "pressure": self.solver.pressure,
            "velocity": self.solver.velocity
        }
        
        # Cell data can also be added if calculated (e.g., vorticity)
        cell_data = {
            # "vorticity": self._calculate_vorticity() # Example
        }

        try:
            self.mesh.write(file_path, point_data=point_data, cell_data=None)
            return file_path
        except Exception as e:
            logging.error(f"Failed to write output file: {e}")
            return None

    def _calculate_vorticity(self) -> np.ndarray:
        """
        (Placeholder) Calculates vorticity at cell centers.
        
        This would require calculating the curl of the velocity field.
        This is a non-trivial calculation on an unstructured mesh.
        """
        logging.debug("Calculating vorticity (placeholder)...")
        # Placeholder: return a zero vector for each cell block
        vorticity_data = {}
        for cell_type, cells in self.mesh.cells.items():
            vorticity_data[cell_type] = np.zeros((len(cells), 3))
        return vorticity_data 